import os
import pandas as pd

class CsvUtil:
    def __init__(self, origin_folder_path, result_folder_path):
        self.origin_folder_path = origin_folder_path
        self.result_folder_path = result_folder_path
        # 确保结果目录存在
        os.makedirs(self.result_folder_path, exist_ok=True)
        # 获取所有csv文件列表
        self.csv_files = [file for file in os.listdir(origin_folder_path) if file.endswith('.csv')]

    def merge_csv(self, merged_file_name="merged"):
        """
        合并指定文件夹中的所有CSV文件[6,8](@ref)
        
        参数:
            merged_file_name: 合并后的文件名（不含扩展名）
        返回:
            output_file: 合并后的文件完整路径
        """
        if not self.csv_files:
            print("警告: 原文件夹中没有找到CSV文件。")
            return None

        # 读取第一个文件作为初始DataFrame
        first_file_path = os.path.join(self.origin_folder_path, self.csv_files[0])
        merged_data = pd.read_csv(first_file_path)
        
        # 循环读取并追加剩余的CSV文件，跳过表头
        for file in self.csv_files[1:]:
            file_path = os.path.join(self.origin_folder_path, file)
            data_to_append = pd.read_csv(file_path)
            merged_data = pd.concat([merged_data, data_to_append], ignore_index=True)
        
        # 构建输出文件路径并保存
        output_file = os.path.join(self.result_folder_path, f"{merged_file_name}.csv")
        merged_data.to_csv(output_file, index=False)
        print(f"合并完成，共合并了 {len(self.csv_files)} 个文件。")
        return output_file

    def split_csv_by_size(self, file_name, chunk_size, output_prefix="split"):
        """
        将指定的CSV文件按固定行数拆分成多个小文件[9,11](@ref)
        
        参数:
            file_name: 待拆分的文件名（需存在于origin_folder_path中）
            chunk_size: 每个小文件包含的行数
            output_prefix: 拆分后文件的前缀
        返回:
            split_num: 拆分后生成的文件数量
        """
        file_path = os.path.join(self.origin_folder_path, file_name)
        if not os.path.exists(file_path):
            print(f"错误: 文件 {file_path} 不存在。")
            return 0

        # 读取原始数据
        data = pd.read_csv(file_path)
        total_rows = len(data)
        
        # 计算需要拆分成多少份
        split_num = (total_rows + chunk_size - 1) // chunk_size  # 向上取整

        for i in range(split_num):
            start_index = i * chunk_size
            end_index = min((i + 1) * chunk_size, total_rows)
            chunk_data = data.iloc[start_index:end_index]
            
            # 生成输出文件名，序号从1开始
            output_file_name = f"{output_prefix}_{i+1}.csv"
            output_file_path = os.path.join(self.result_folder_path, output_file_name)
            
            chunk_data.to_csv(output_file_path, index=False)
        
        print(f"拆分完成，共生成 {split_num} 个文件，每个文件最多 {chunk_size} 行。")
        return split_num

    def split_csv_by_count(self, file_name, num_pieces, output_prefix="split"):
        """
        将指定的CSV文件拆分成指定数量的小文件[9](@ref)
        
        参数:
            file_name: 待拆分的文件名（需存在于origin_folder_path中）
            num_pieces: 想要拆分成多少个小文件
            output_prefix: 拆分后文件的前缀
        返回:
            actual_pieces: 实际拆分生成的文件数量
        """
        file_path = os.path.join(self.origin_folder_path, file_name)
        if not os.path.exists(file_path):
            print(f"错误: 文件 {file_path} 不存在。")
            return 0

        # 读取原始数据
        data = pd.read_csv(file_path)
        total_rows = len(data)
        
        # 计算每个文件大致的行数
        base_chunk_size = total_rows // num_pieces
        remainder = total_rows % num_pieces

        start_index = 0
        actual_pieces = 0
        for i in range(num_pieces):
            # 计算当前分块的行数：前面的分块多承担一些余数
            current_chunk_size = base_chunk_size + (1 if i < remainder else 0)
            if current_chunk_size == 0:
                break  # 避免生成空文件

            end_index = start_index + current_chunk_size
            chunk_data = data.iloc[start_index:end_index]
            
            output_file_name = f"{output_prefix}_{i+1}.csv"
            output_file_path = os.path.join(self.result_folder_path, output_file_name)
            
            chunk_data.to_csv(output_file_path, index=False)
            
            start_index = end_index
            actual_pieces += 1
        
        print(f"拆分完成，共生成 {actual_pieces} 个文件。")
        return actual_pieces


if __name__ == '__main__':
    # 示例用法
    origin_folder_path = r'learn-python-easily/04_Python_Practice/P01_CSVUtil/DataSet'  # 请更换为你的原始CSV文件夹路径
    result_folder_path = r'learn-python-easily/04_Python_Practice/P01_CSVUtil/newDataSet'  # 请更换为你想要保存结果的文件夹路径

    csv_util = CsvUtil(origin_folder_path, result_folder_path)

    # 示例1: 合并文件
    merged_file_path = csv_util.merge_csv("merged_result")
    if merged_file_path:
        print(f"合并后的文件保存在：{merged_file_path}")

    # 示例2: 按指定行数拆分某个文件 (例如拆分 merged_result.csv, 每1000行一个文件)
    # split_count_by_size = csv_util.split_csv_by_size("merged_result.csv", 1000, "split_by_size")
    # print(f"按行数拆分，共生成了 {split_count_by_size} 个文件。")

    # 示例3: 按指定份数拆分某个文件 (例如将 merged_result.csv 拆分成5份)
    # split_count_by_num = csv_util.split_csv_by_count("merged_result.csv", 5, "split_by_count")
    # print(f"按份数拆分，共生成了 {split_count_by_num} 个文件。")